# -*- coding: utf-8 -*-
from scapy.all import *
import os
 
def read_pcap(file_name):
    print('开始读取pcap')
    file_data = rdpcap(file_name)
    print('读取完成')
    return file_data
def fenliu(file_data):
    port = []
    l = len(file_data)
    for i in range(0,l):
        try:
            # syn包的tcp标志位为0x02,跟据syn包先将源端口组成列表，返回port列表
            if file_data[i]['TCP'].flags == 0x02:
                port.append(file_data[i]['TCP'].sport)
        except:
            pass
    return port
def fenbao(file_data):
    port = fenliu(file_data)
    l1 = len(port)
    data_list=[[] for _ in range(l1)]
    #data_list=[[]]*l1  不能这样创建多维列表，会出问题，应该用上面的方式
    l = len(file_data)
    for i in range(0, l):
        try:
            #源端口在port列表里，则端口获取在列表中的索引，放入对应的data_list中
            if file_data[i]['TCP'].sport in port :
                j = port.index(file_data[i]['TCP'].sport)
                data_list[j].append(file_data[i])
            # 目的端口在port列表里，则端口获取在列表中的索引，放入对应的data_list中
            elif file_data[i]['TCP'].dport in port:
                j = port.index(file_data[i]['TCP'].dport)
                data_list[j].append(file_data[i])
        except:
            pass
    new_data_list = [ele for ele in data_list if ele !=[]]
    # for i in range(len(new_data_list)):
    #     print(new_data_list[i])
    return new_data_list
def save_file(dir, data_list):
    l = len(data_list)
    print('开始保存pcap,共有'+str(l)+'个pcap包')
    for i in range(0,l):
        name = dir + 'ips'+str(i)+'.pcap'
        wrpcap(name, data_list[i])
    print('保存完成')
if __name__ =="__main__":
    dir_name = os.path.dirname(__file__) + '/'
    file_name = dir_name + "ss5_3-4-4.pcap"
    # file_name = dir_name + "baidu2.pcap"
    file_data = read_pcap(file_name)
    data=fenbao(file_data)
    save_file(dir_name,data)